# -*- coding:utf-8 -*-
# @Author: shenyuyu
# @Time: 2023/6/29 17:29
# @File: qu_2_自定义udf函数.py
from pyspark.sql import SparkSession
from pyspark.sql.types import IntegerType
from pyspark.sql import functions as F

if __name__ == '__main__':
    spark = SparkSession.builder.appName("a").master("local[*]").getOrCreate()
    sc = spark.sparkContext
    rdd = sc.parallelize([1, 2, 3, 4, 5]).map(lambda x: [x])
    df = rdd.toDF(["num"])
    df.createOrReplaceTempView("nums")


    def udf_fun(data):
        return data * 10

    udf1 = spark.udf.register("udf2", udf_fun, IntegerType())

    spark.sql("select udf2(num) from nums").show()

    df.selectExpr("udf2(num)")

    df.select(udf1("num")).withColumnRenamed("udf2(num)", "num10").show()

    udf3 = F.udf(udf_fun, IntegerType())

    df.select(udf3("num")).withColumnRenamed("udf_fun(num)", "num10").show()